import cv2 as cv
import argparse
import sys
import numpy as np
import os.path

# 初始化参数
confThreshold = 0.4  # 置信阈值
nmsThreshold = 0.3  # 非最大抑制阈值
inpWidth = 608  # 网络输入图像的宽度320 416 608
inpHeight = 608  # 高度

parser = argparse.ArgumentParser(description='Object Detection using YOLO in OPENCV')
parser.add_argument('--device', default='cpu', help="Device to perform inference on 'cpu' or 'gpu'.")
parser.add_argument('--image', help='Path to image file.')
parser.add_argument('--video', help='Path to video file.')
args = parser.parse_args()

# 加载classes
classesFile = "person.names"
classes = None
with open(classesFile, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')

# 给出模型的配置和权重文件，并使用它们加载网络。
modelConfiguration = "person.cfg"
modelWeights = "person_32_2000.weights"

net = cv.dnn.readNetFromDarknet(modelConfiguration, modelWeights)

if (args.device == 'cpu'):
    net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
    net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)
    print('Using CPU device.')
elif (args.device == 'gpu'):
    net.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
    net.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
    print('Using GPU device.')


# 获取输出层的名称
def getOutputsNames(net):
    # 获取网络中所有层的名称
    layersNames = net.getLayerNames()
    # 获取输出层的名称，即具有未连接输出的层
    return [layersNames[i - 1] for i in net.getUnconnectedOutLayers()]


# 绘制预测的边界框
def drawPred(classId, conf, left, top, right, bottom):
    # 绘制边界框
    cv.rectangle(frame, (left, top), (right, bottom), (255, 178, 50), 3)

    label = '%.2f' % conf

    # 获取类名及其置信度的标签
    if classes:
        assert (classId < len(classes))
        label = '%s:%s' % (classes[classId], label)

    # 在边界框顶部显示标签
    labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
    top = max(top, labelSize[1])
    cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine),
                 (255, 255, 255), cv.FILLED)
    cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 0), 1)


# 使用非最大值抑制移除低置信度的边界框
def postprocess(frame, outs):
    frameHeight = frame.shape[0]
    frameWidth = frame.shape[1]

    # 扫描网络输出的所有边界框，仅保留置信高的人。将框的类标签指定为得分最高的类。
    classIds = []
    confidences = []
    boxes = []
    for out in outs:
        for detection in out:
            scores = detection[5:]
            classId = np.argmax(scores)
            if (classId == 0):  # 只检测第一个class
            	confidence = scores[classId]
            	if confidence > confThreshold:
					
                    center_x = int(detection[0] * frameWidth)
                    center_y = int(detection[1] * frameHeight)
                    width = int(detection[2] * frameWidth)
                    height = int(detection[3] * frameHeight)
                    left = int(center_x - width / 2)
                    top = int(center_y - height / 2)
                    classIds.append(classId)
                    confidences.append(float(confidence))
                    boxes.append([left, top, width, height])

    # 执行非最大抑制以消除置信度较低的冗余重叠框。
    indices = cv.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold)
    for i in indices:
        # i = i[0]
        box = boxes[i]
        left = box[0]
        top = box[1]
        width = box[2]
        height = box[3]
        drawPred(classIds[i], confidences[i], left, top, left + width, top + height)


# 过程输入
winName = 'Deep learning object detection in OpenCV'
cv.namedWindow(winName, cv.WINDOW_NORMAL)

outputFile = "yolo_out_person_.avi"

if (args.image):
    # 打开图像文件
    if not os.path.isfile(args.image):
        print("Input image file ", args.image, " doesn't exist")
        sys.exit(1)
    cap = cv.VideoCapture(args.image)
    outputFile = args.image[:-4] + '_yolo_out_.jpg'
elif (args.video):
    # 打开视频文件
    if not os.path.isfile(args.video):
        print("Input video file ", args.video, " doesn't exist")
        sys.exit(1)
    cap = cv.VideoCapture(args.video)
    outputFile = args.video[:-4] + '_yolo_out_person.avi'
else:
    # 网络摄像头输入
    cap = cv.VideoCapture(0)

# 初始化视频写入程序以保存输出视频
if (not args.image):
    vid_writer = cv.VideoWriter(outputFile, cv.VideoWriter_fourcc('M', 'J', 'P', 'G'), 30,
                                (round(cap.get(cv.CAP_PROP_FRAME_WIDTH)), round(cap.get(cv.CAP_PROP_FRAME_HEIGHT))))

while cv.waitKey(1) < 0:

    # 从视频中获取帧
    hasFrame, frame = cap.read()

    # 如果到达视频末尾，请停止程序
    if not hasFrame:
        print("Done processing !!!")
        print("Output file is stored as ", outputFile)
        cv.waitKey(3000)
        # Release device
        cap.release()
        break

    # 对图像进行预处理，包括减均值，比例缩放，裁剪，交换通道等，返回一个4通道的blob(blob可以简单理解为一个N维的数组，用于神经网络的输入)
    blob = cv.dnn.blobFromImage(frame, 1 / 255, (inpWidth, inpHeight), [0, 0, 0], 1, crop=False)

    # 设置网络输入
    net.setInput(blob)

    # 运行正向传递以获得输出层的输出
    outs = net.forward(getOutputsNames(net))

    # 以低置信度删除边界框
    postprocess(frame, outs)

    # 输入效率信息，函数getPerfProfile返回推理的总时间（t）和每个层的计时（以层时间为单位）
    t, _ = net.getPerfProfile()
    label = 'Inference time: %.2f ms' % (t * 1000.0 / cv.getTickFrequency())
    cv.putText(frame, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))

    # 用检测框书写框架
    if (args.image):
        cv.imwrite(outputFile, frame.astype(np.uint8))
    else:
        vid_writer.write(frame.astype(np.uint8))

    cv.imshow(winName, frame)
